package cn.zh.aop.component;

import cn.zh.aop.constant.DataSourceEnum;
import lombok.RequiredArgsConstructor;
import org.springframework.context.annotation.Primary;
import org.springframework.context.annotation.Profile;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.stereotype.Component;

import javax.sql.DataSource;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author zh
 */
@Primary
@Component
@RequiredArgsConstructor
@Profile("dev")
public class DataSourceManagement extends AbstractRoutingDataSource {

    /**
     * 这种方式每次进来都会对flag初始化
     */
//    public static ThreadLocal<String> flag = ThreadLocal.withInitial(DataSourceManagement::flagInit);

    /**
     * 这种方式每次进来不会对flag初始化，但是由于配置了默认数据源，所以当 determineCurrentLookupKey() 返回空，则使用默认数据源
     */
    public static ThreadLocal<String> flag = new ThreadLocal<>();

    private final DataSource masterDataSource;

    private final DataSource slaveDataSource;

    public static String flagInit(){
        return DataSourceEnum.MASTER.name();
    }

    @Override
    protected Object determineCurrentLookupKey() {
        return flag.get();
    }

    @Override
    public void afterPropertiesSet() {
        Map<Object,Object> targetDataSource = new ConcurrentHashMap<>();
        targetDataSource.put(DataSourceEnum.MASTER.name(),masterDataSource);
        targetDataSource.put(DataSourceEnum.SLAVE.name(),slaveDataSource);
        // 设置多种数据源
        super.setTargetDataSources(targetDataSource);
        // 设置默认数据源
        super.setDefaultTargetDataSource(masterDataSource);
        super.afterPropertiesSet();

    }
}
